import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure


class ImageViewer(QLabel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setAlignment(Qt.AlignCenter)
        self.setText("拖拽图像到此处或使用菜单打开")
        self.setMinimumSize(640, 480)
        self.original_image = None

    def load_image(self, file_path):
        self.original_image = cv2.imread(file_path)
        if self.original_image is None:
            QMessageBox.warning(self, "错误", "无法加载图像!")
            return False

        self.show_image(self.original_image)
        return True

    def show_image(self, image):
        if image is None:
            return

        # 转换OpenCV BGR图像为Qt RGB格式
        h, w, c = image.shape
        bytes_per_line = 3 * w
        q_img = QImage(image.data, w, h, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
        pixmap = QPixmap.fromImage(q_img)
        self.setPixmap(pixmap.scaled(self.width(), self.height(), Qt.KeepAspectRatio))


class CVPlatform(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("计算机视觉开发平台")
        self.setGeometry(100, 100, 1200, 800)

        # 创建菜单栏
        self.create_menus()

        # 主布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QHBoxLayout()
        main_widget.setLayout(main_layout)

        # 左侧控制面板
        control_panel = QGroupBox("图像处理控制")
        control_layout = QVBoxLayout()
        control_panel.setLayout(control_layout)
        control_panel.setFixedWidth(300)

        # 图像处理功能
        self.create_image_controls(control_layout)
        control_layout.addStretch()

        # 形状识别功能
        self.create_shape_controls(control_layout)
        control_layout.addStretch()

        # 视频处理功能
        self.create_video_controls(control_layout)

        # 图像显示区域
        self.image_viewer = ImageViewer()

        # 添加到主布局
        main_layout.addWidget(control_panel)
        main_layout.addWidget(self.image_viewer)

        # 状态栏
        self.statusBar().showMessage("就绪")

    def create_menus(self):
        menubar = self.menuBar()

        # 文件菜单
        file_menu = menubar.addMenu('文件')
        open_action = QAction('打开图像', self)
        open_action.triggered.connect(self.open_image)
        file_menu.addAction(open_action)

        # 编辑菜单
        edit_menu = menubar.addMenu('编辑')
        rotate_action = QAction('旋转90度', self)
        rotate_action.triggered.connect(lambda: self.rotate_image(90))
        edit_menu.addAction(rotate_action)

        # 视图菜单
        view_menu = menubar.addMenu('视图')
        reset_action = QAction('重置视图', self)
        reset_action.triggered.connect(self.reset_view)
        view_menu.addAction(reset_action)

    def create_image_controls(self, layout):
        group = QGroupBox("基础图像处理")
        form = QFormLayout()

        # 亮度调整
        self.brightness_slider = QSlider(Qt.Horizontal)
        self.brightness_slider.setRange(-100, 100)
        self.brightness_slider.setValue(0)
        self.brightness_slider.valueChanged.connect(self.adjust_brightness_contrast)
        form.addRow("亮度:", self.brightness_slider)

        # 对比度调整
        self.contrast_slider = QSlider(Qt.Horizontal)
        self.contrast_slider.setRange(-100, 100)
        self.contrast_slider.setValue(0)
        self.contrast_slider.valueChanged.connect(self.adjust_brightness_contrast)
        form.addRow("对比度:", self.contrast_slider)

        # 模糊处理
        self.blur_slider = QSlider(Qt.Horizontal)
        self.blur_slider.setRange(0, 30)
        self.blur_slider.setValue(0)
        self.blur_slider.valueChanged.connect(self.apply_blur)
        form.addRow("模糊强度:", self.blur_slider)

        # 边缘检测按钮
        edge_btn = QPushButton("边缘检测(Canny)")
        edge_btn.clicked.connect(self.detect_edges)
        form.addRow(edge_btn)

        # 重置按钮
        reset_btn = QPushButton("重置调整")
        reset_btn.clicked.connect(self.reset_adjustments)
        form.addRow(reset_btn)

        group.setLayout(form)
        layout.addWidget(group)

    def create_shape_controls(self, layout):
        group = QGroupBox("形状识别")
        vbox = QVBoxLayout()

        detect_btn = QPushButton("检测形状")
        detect_btn.clicked.connect(self.detect_shapes)
        vbox.addWidget(detect_btn)

        self.shape_result = QLabel("检测结果将显示在这里")
        vbox.addWidget(self.shape_result)

        group.setLayout(vbox)
        layout.addWidget(group)

    def create_video_controls(self, layout):
        group = QGroupBox("视频处理")
        vbox = QVBoxLayout()

        self.webcam_btn = QPushButton("开启摄像头")
        self.webcam_btn.setCheckable(True)
        self.webcam_btn.toggled.connect(self.toggle_webcam)
        vbox.addWidget(self.webcam_btn)

        self.video_label = QLabel("摄像头将显示在这里")
        self.video_label.setAlignment(Qt.AlignCenter)
        self.video_label.setMinimumSize(320, 240)
        vbox.addWidget(self.video_label)

        group.setLayout(vbox)
        layout.addWidget(group)

        # 视频捕获对象
        self.cap = None
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_frame)

    def open_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "打开图像", "",
            "图像文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if file_path:
            if self.image_viewer.load_image(file_path):
                self.statusBar().showMessage(f"已加载: {file_path}")

    def rotate_image(self, angle):
        if self.image_viewer.original_image is None:
            return

        # 获取中心点并旋转
        image = self.image_viewer.original_image
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(image, M, (w, h))
        self.image_viewer.show_image(rotated)

    def adjust_brightness_contrast(self):
        if self.image_viewer.original_image is None:
            return

        brightness = self.brightness_slider.value()
        contrast = self.contrast_slider.value()

        # 应用亮度和对比度调整
        alpha = 1 + contrast / 100
        beta = brightness
        adjusted = cv2.convertScaleAbs(self.image_viewer.original_image, alpha=alpha, beta=beta)
        self.image_viewer.show_image(adjusted)

    def apply_blur(self):
        if self.image_viewer.original_image is None:
            return

        ksize = self.blur_slider.value() * 2 + 1  # 确保是奇数
        if ksize <= 0:
            self.image_viewer.show_image(self.image_viewer.original_image)
            return

        blurred = cv2.GaussianBlur(self.image_viewer.original_image, (ksize, ksize), 0)
        self.image_viewer.show_image(blurred)

    def detect_edges(self):
        if self.image_viewer.original_image is None:
            return

        gray = cv2.cvtColor(self.image_viewer.original_image, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        # 转换为彩色显示
        edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
        self.image_viewer.show_image(edges_colored)

    def detect_shapes(self):
        if self.image_viewer.original_image is None:
            return

        # 转换为灰度图并检测轮廓
        gray = cv2.cvtColor(self.image_viewer.original_image, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        _, thresh = cv2.threshold(blurred, 60, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 复制图像用于绘制
        image_copy = self.image_viewer.original_image.copy()
        shapes = []

        for contour in contours:
            # 忽略小轮廓
            if cv2.contourArea(contour) < 500:
                continue

            # 多边形逼近
            peri = cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, 0.04 * peri, True)
            sides = len(approx)

            # 根据边数判断形状
            shape = ""
            if sides == 3:
                shape = "三角形"
            elif sides == 4:
                # 检查是否为矩形
                x, y, w, h = cv2.boundingRect(approx)
                aspect_ratio = w / float(h)
                shape = "正方形" if 0.95 <= aspect_ratio <= 1.05 else "矩形"
            elif sides > 4:
                shape = f"{sides}边形"
            else:
                # 可能是圆形
                (x, y), radius = cv2.minEnclosingCircle(contour)
                shape = "圆形"

            if shape:
                shapes.append(shape)
                # 绘制轮廓和标签
                cv2.drawContours(image_copy, [contour], -1, (0, 255, 0), 2)
                M = cv2.moments(contour)
                if M["m00"] != 0:
                    cX = int(M["m10"] / M["m00"])
                    cY = int(M["m01"] / M["m00"])
                    cv2.putText(image_copy, shape, (cX, cY),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

        self.image_viewer.show_image(image_copy)
        self.shape_result.setText(f"检测到形状: {', '.join(shapes)}" if shapes else "未检测到形状")

    def reset_adjustments(self):
        self.brightness_slider.setValue(0)
        self.contrast_slider.setValue(0)
        self.blur_slider.setValue(0)
        if self.image_viewer.original_image is not None:
            self.image_viewer.show_image(self.image_viewer.original_image)

    def reset_view(self):
        if self.image_viewer.original_image is not None:
            self.image_viewer.show_image(self.image_viewer.original_image)
        self.reset_adjustments()
        self.shape_result.setText("检测结果将显示在这里")

    def toggle_webcam(self, checked):
        if checked:
            self.cap = cv2.VideoCapture(0)
            if not self.cap.isOpened():
                QMessageBox.warning(self, "错误", "无法访问摄像头!")
                self.webcam_btn.setChecked(False)
                return
            self.timer.start(30)  # 30ms更新一次
            self.webcam_btn.setText("关闭摄像头")
        else:
            self.timer.stop()
            if self.cap:
                self.cap.release()
            self.video_label.clear()
            self.video_label.setText("摄像头已关闭")
            self.webcam_btn.setText("开启摄像头")

    def update_frame(self):
        ret, frame = self.cap.read()
        if ret:
            # 转换为RGB格式
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, c = frame.shape
            bytes_per_line = c * w
            q_img = QImage(frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(q_img)
            self.video_label.setPixmap(pixmap.scaled(
                self.video_label.width(),
                self.video_label.height(),
                Qt.KeepAspectRatio
            ))


if __name__ == "__main__":
    app = QApplication(sys.argv)
    platform = CVPlatform()
    platform.show()
    sys.exit(app.exec_())